import os
import sys

import sys 
sys.path.append("..") 
sys.path.append("./") 
sys.path.append("./model")

import torch
import argparse
import logging
import sklearn
from os.path import join
from datetime import datetime
from torch.utils.model_zoo import load_url
from google_drive_downloader import GoogleDriveDownloader as gdd
import cv2

# import test_3d
# import util
import commons
import datasets.datasets_ws as datasets_ws
from model import network
import trimesh
import numpy as np

from model.extractors.FeatureVectorExtractor import FeatureVectorExtractor

# from model.common.Context import Context

from tqdm import tqdm
import yaml

from model.search.LocalFeatureSet2 import LocalFeatureSet2
from torch.utils.data.dataloader import DataLoader
from model.utils.htimer import htimer

import model.utils.parse_args as parse_args

context = argparse.Namespace()
args,config = parse_args.parse_arguments("./configs/test_3dtiles_3d.yaml")
context.args = args
context.config = config

# print("check config!! ", config)

start_time = datetime.now()
context.save_dir = join("log", args.save_dir, start_time.strftime('%Y-%m-%d_%H-%M-%S'))
commons.setup_logging(context.save_dir)
# commons.make_deterministic(args.seed)
logging.info(f"Arguments: {args}")
logging.info(f"The outputs are being saved in {context.save_dir}")



###########################################FE###########################################
model = FeatureVectorExtractor(args)
model = model.to(args.device)


# Enable DataParallel after loading checkpoint, otherwise doing it before
# would append "module." in front of the keys of the state dict triggering errors
model = torch.nn.DataParallel(model)


####################################

from datasets.dataset_3dtiles import Dataset3DTiles

if args.data.dataset_type == "3DTiles":
    test_ds = Dataset3DTiles(context)

context.dataset = test_ds
context.fve_model = model
# context.test_method = args.test_method


class RTreeIntersection():
    def __init__(self) -> None:
        self._bounds = None
        self._vectors = None
        self._rtree = None
        self._triangles = None
        self._uv = None
        pass

    def build(self,mesh,w,h):
        uvAll = mesh.visual.uv
        facesAll = mesh.faces
        # facesNormalAll = gltf_mesh.face_normals
        vertices = mesh.vertices

        uvAll[:,0] *= w
        uvAll[:,1] = h*(1-uvAll[:,1])

        uv2 = uvAll[facesAll]

        uv_max = np.max(uv2,axis=1)
        uv_min = np.min(uv2,axis=1)

        uv_b = np.concatenate((uv_min,uv_max),axis=1)
        rtree = trimesh.util.bounds_tree(uv_b)
        self._rtree  = rtree
        self._triangles = mesh.triangles
        self._uv = uv2



    
    def insert(self,P):
        def _inface(P,uv3):
          
            A = uv3[0]
            B = uv3[1]
            C = uv3[2]
            v0 = C - A 
            v1 = B - A 
            v2 = P - A 
            dot00 = np.dot(v0,v0) 
            dot01 = np.dot(v0,v1) 
            dot02 = np.dot(v0,v2) 
            dot11 = np.dot(v1,v1) 
            dot12 = np.dot(v1,v2) 
            inverDeno = 1 / (dot00 * dot11 - dot01 * dot01) ;
            u = (dot11 * dot02 - dot01 * dot12) * inverDeno
            # if u < 0 or u > 1 :
            #     return False
            v = (dot00 * dot12 - dot01 * dot02) * inverDeno 
            # if (v < 0 or v > 1):
            #     return False
            return u,v
        
        pp33 = None

        res_list = list(self._rtree.intersection(P) )
        for t_ind in res_list:
            tri = self._uv[t_ind]
            nu,nv = _inface(P,tri)
            if (nu < 0 or nu > 1) or (nv < 0 or nv > 1) or (nu + nv >1):
                continue
            
            p3 = self._triangles[t_ind]
            pA = p3[0]
            pB = p3[1]
            pC = p3[2]
            pu = pC - pA
            pv = pB - pA

            pp33 = pA + pu * nu + pv * nv
            break
        return pp33

class TriTop:
    def __init__(self) -> None:
        self._faces = None
        self._uv = None
        self._points = None
        pass


    
    def insert(self,P):
        def _inface(P,uv3):
            # if P[0] > np.max(uv3[0:]):
            #     return 2,0
            # if P[0] < np.min(uv3[0:]):
            #     return -2,0
            # if P[1] > np.max(uv3[:1]):
            #     return 0,2
            # if P[1] < np.min(uv3[:1]):
            #     return 0,-2
            
            A = uv3[0]
            B = uv3[1]
            C = uv3[2]
            v0 = C - A 
            v1 = B - A 
            v2 = P - A 
            dot00 = np.dot(v0,v0) 
            dot01 = np.dot(v0,v1) 
            dot02 = np.dot(v0,v2) 
            dot11 = np.dot(v1,v1) 
            dot12 = np.dot(v1,v2) 
            inverDeno = 1 / (dot00 * dot11 - dot01 * dot01) ;
            u = (dot11 * dot02 - dot01 * dot12) * inverDeno
            # if u < 0 or u > 1 :
            #     return False
            v = (dot00 * dot12 - dot01 * dot02) * inverDeno 
            # if (v < 0 or v > 1):
            #     return False
            return u,v
        
        pp33 = None
        for i in range(self._faces.shape[0]):
            face = self._faces[i]
            uv3 = self._uv[face]
            nu,nv = _inface(P,uv3)
            if (nu < 0 or nu > 1) or (nv < 0 or nv > 1) or (nu + nv >1):
                continue
            
            p3 = self._points[face]
            pA = p3[0]
            pB = p3[1]
            pC = p3[2]
            pu = pC - pA
            pv = pB - pA

            pp33 = pA + pu * nu + pv * nv
            break
        return pp33


#计算features.h5中，每个特征点的三维坐标
def _getPoint3DCoordinates(args,context):
    feature_set = LocalFeatureSet2()
    fdir = join(args.data.datasets_folder,args.data.dataset_name,args.features.features_dir)
    feature_set.set_feature_dir(fdir)
    if feature_set.read() < 0:
        return -1
    
    features = feature_set.get_features()
    
    htimer.set_qs_type("INSERT")

    pbar = tqdm(total=len(features))

    for id,fd in features.items():
        pbar.update()
        fname = fd["image_name"]
        glb_fname =  fname.replace(".jpg","")
        glb_fname =  join(args.data.datasets_folder,args.data.dataset_name,glb_fname)
        bottle_gltf = trimesh.load(glb_fname)
        gltf_mesh = bottle_gltf.geometry['GLTF']
        cvImg = None
        pilImg = gltf_mesh.visual.material.baseColorTexture
        if pilImg is None:
            continue

        cvImg = cv2.cvtColor(np.array(pilImg),cv2.COLOR_RGB2BGR)

        htimer.start_time()

        width = cvImg.shape[1]
        height =cvImg.shape[0]


        #######快速版本
        tritop = RTreeIntersection()
        tritop.build(gltf_mesh,width,height)
        
        #######慢速版本
        # uvAll = gltf_mesh.visual.uv
        # facesAll = gltf_mesh.faces
        # facesNormalAll = gltf_mesh.face_normals
        # vertices = gltf_mesh.vertices

        # uvAll[:,0] *= width
        # uvAll[:,1] = height*(1-uvAll[:,1])       

        # tritop = TriTop()
        # tritop._faces = facesAll    
        # tritop._uv = uvAll
        # tritop._points = vertices

        keypoints = fd["keypoints"]
        
        points3ds = np.zeros((keypoints.shape[0],3),dtype=keypoints.dtype)
        
        for i in range(keypoints.shape[0]):
            p = keypoints[i]
            p = p[0:2]

            c3 = tritop.insert(p)
            if c3 is not None:
                points3ds[i]=c3
        
        fd["points3ds"] = points3ds
        feature_set.update_features(id,"points3ds",points3ds)
        # feature_set.update_features(id,"attrs/image_width",width)
        # feature_set.update_features(id,"attrs/image_height",height)
        htimer.end_time()

    pbar.close()
    print(htimer.gettime())
    return 0

_getPoint3DCoordinates(args,context)